package com.leyantech.ecp.agenx.connector.handler;

import com.leyantech.ecp.agenx.common.config.Constants;
import com.leyantech.ecp.agenx.common.config.MessageAddress;
import com.leyantech.ecp.agenx.common.misc.TracingHelper;
import com.leyantech.ecp.agenx.common.sender.Sender;
import com.leyantech.utility.config.ConfigUtil;

import com.ctrip.framework.apollo.Config;
import com.ctrip.framework.apollo.ConfigService;
import com.google.common.util.concurrent.RateLimiter;
import com.google.inject.Inject;
import com.google.inject.assistedinject.Assisted;
import io.opentracing.Scope;
import io.opentracing.Span;
import io.opentracing.Tracer;
import io.opentracing.util.GlobalTracer;
import io.vertx.core.Handler;
import io.vertx.core.MultiMap;
import io.vertx.core.Promise;
import io.vertx.core.buffer.Buffer;
import io.vertx.core.eventbus.EventBus;
import io.vertx.core.eventbus.Message;
import io.vertx.core.eventbus.MessageConsumer;
import io.vertx.core.http.ServerWebSocket;
import io.vertx.core.http.WebSocketFrame;
import io.vertx.core.impl.VertxInternal;
import io.vertx.core.json.JsonObject;
import io.vertx.ext.auth.User;
import io.vertx.tracing.opentracing.OpenTracingUtil;
import lombok.extern.slf4j.Slf4j;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.slf4j.MDC;

import java.util.concurrent.TimeUnit;
import java.util.function.Supplier;

/**
 * websocket 消息处理器，注意：eventLoop 线程一般和 cpu 核心数差不多，所以耗时的任务（发送 kafka，关闭连接等 网络操作）都交给 workPool
 *
 * @author eahau, {@literal <eahau@foxmail.com>}
 * @date 2022-10-27.
 */
@Slf4j
public class WebSocketMessageHandler implements Handler<ServerWebSocket> {

  private static final Logger inLogger = LoggerFactory.getLogger("com.leyantech.ecp.agenx.connector.WebSocket.In");

  private static final Logger outLogger = LoggerFactory.getLogger("com.leyantech.ecp.agenx.connector.WebSocket.Out");

  static final Config RATE_LIMITER_CONFIG = ConfigService.getConfig("rateLimiter");

  private final VertxInternal vertx;

  private final EventBus eventBus;

  private final RateLimiter rateLimiter;

  private final String storeId;

  private final String assistantId;

  private static final Supplier<Long> heartbeatStepSeconds = ConfigUtil.getDynamicConfig(
      "heartbeat.step.seconds",
      seconds -> TimeUnit.SECONDS.toMillis(Long.parseLong(seconds)),
      TimeUnit.SECONDS.toMillis(10L)
  );

  private static final Supplier<Integer> heartbeatCheckCount = ConfigUtil.getDynamicConfig(
      "heartbeat.check.count",
      Integer::parseInt,
      3
  );

  final String connectorId;

  private volatile long lastActiveTimestamp = System.currentTimeMillis();

  private final Promise<Void> registerCompleteHandler;

  private final Sender<String> messageSender;

  private final JsonObject token;

  @Inject
  public WebSocketMessageHandler(VertxInternal vertx,
      @Assisted User user,
      @Assisted Promise<Void> registerCompleteHandler,
      Sender<String> messageSender) {
    this.vertx = vertx;
    this.eventBus = vertx.eventBus();
    this.token = user.attributes().getJsonObject("accessToken");
    this.storeId = user.get(Constants.STORE_ID_FIELD);
    this.assistantId = user.get(Constants.ASSISTANT_ID_FIELD);
    this.messageSender = messageSender;
    this.connectorId = String.join(".", storeId, assistantId);
    this.registerCompleteHandler = registerCompleteHandler;
    token.put("node_id", vertx.getClusterManager().getNodeId());
    this.rateLimiter = getRateLimiter();
  }

  private RateLimiter getRateLimiter() {
    final Config config = RATE_LIMITER_CONFIG;

    Double rate = config.getDoubleProperty(connectorId, null);
    if (rate == null) {
      rate = config.getDoubleProperty(storeId, 100.d);
    }

    return RateLimiter.create(rate);
  }

  @Override
  public void handle(ServerWebSocket websocket) {

    final Handler<Message<String>> responseHandler = message -> {
      final Span span = OpenTracingUtil.getSpan();
      vertx.executeBlocking(promise -> {
        try (final Scope ignore = GlobalTracer.get().scopeManager().activate(span)) {
          mdcSetup();
          final String responseMessage = message.body();
          final MultiMap headers = message.headers();

          outLogger.info("received reply, message {}", responseMessage);

          // 内部消息
          if ("cmd".equals(headers.get("type"))) {
            final JsonObject command = new JsonObject(responseMessage);
            final JsonObject header = command.getJsonObject("header");
            final String fn;
            if (header == null || (fn = header.getString("fn")) == null) {
              return;
            }
            outLogger.warn("connector {} received cmd {} {}", connectorId, fn, responseMessage);
            switch (fn) {
              // 踢下线
              case "KickOut":
                final JsonObject body = command.getJsonObject("body", new JsonObject());
                final short code = body.getInteger("code", 1000).shortValue();
                final String reason = body.getString("reason", "KickOut");
                final String handlerID = body.getString("handlerID");
                // 来自当前链接的消息，忽略
                if (!websocket.textHandlerID().equals(handlerID)) {
                  websocket.close(code, reason);
                }
                break;
              // 判断客服是否在线
              case "Ping":
                message.reply(websocket.isClosed() ? "Closed" : "Pong");
                break;
            }
          } else {
            // 返给客户端的消息
            websocket.writeTextMessage(responseMessage)
                .onComplete(TracingHelper.tracingHandler(res -> {
                      if (res.succeeded()) {
                        outLogger.info("send message to client succeed. {}", responseMessage);
                      } else {
                        outLogger
                            .error("send message to client failed. {}", responseMessage, res.cause());
                      }
                    }, () -> span)
                );
          }
          promise.complete();
        } catch (Throwable throwable) {
          promise.fail(throwable);
          outLogger.error("send message to client failed.", throwable);
        } finally {
          MDC.clear();
        }
      });
    };

    final MessageConsumer<String> assistantConsumer = eventBus.consumer(
        MessageAddress.toConnector(connectorId), responseHandler);

    // 店铺级别广播
    final MessageConsumer<String> storeConsumer = eventBus
        .consumer(MessageAddress.toConnector(storeId), responseHandler);
    // 全局广播
    final MessageConsumer<String> allConsumer = eventBus
        .consumer(MessageAddress.ALL_CONNECTOR, responseHandler);

    // 注册到集群，callback
    assistantConsumer.completionHandler(result -> {
      if (result.succeeded()) {
        registerCompleteHandler.complete();
      } else {
        registerCompleteHandler.fail(result.cause());
      }
    });

    final String handlerID = websocket.textHandlerID();

    final long heartbeatTimerId = registerHeartbeat(websocket);

    final Handler<Void> closeHandler = Void -> vertx.executeBlocking(promise -> {
      sendAssistantMessage(false);
      cancelHeartbeat(heartbeatTimerId);

      assistantConsumer.unregister()
          .onFailure(throwable -> log.error("{} assistantConsumer unregister failed.",
              connectorId, throwable));

      storeConsumer.unregister()
          .onFailure(throwable -> log.error("{} storeConsumer unregister failed.",
              connectorId, throwable));

      allConsumer.unregister()
          .onFailure(throwable -> log.error("{} allConsumer unregister failed.",
              connectorId, throwable));

      promise.complete();
      log.info("ws closed {} {} {}", storeId, assistantId, handlerID);
    });

    final Handler<String> requestHandler = message -> vertx.executeBlocking(promise -> {
      lastActiveTimestamp = System.currentTimeMillis();
      final Tracer tracer = GlobalTracer.get();
      final Span span = tracer.buildSpan("ws-received").start();
      try (final Scope ignore = tracer.scopeManager().activate(span)) {
        mdcSetup();
        inLogger.info("received text message {}", message);
        if (!rateLimiter.tryAcquire()) {
          inLogger.warn("received text message too many times");
          return;
        }
        messageSender.send(connectorId, message);
        promise.complete();
      } catch (Throwable throwable) {
        promise.fail(throwable);
        inLogger.error("received text message error, message {}", message, throwable);
      } finally {
        if (span != null) {
          span.finish();
        }
        MDC.clear();
      }
    });

    websocket.textMessageHandler(requestHandler)
        .pongHandler(ping -> lastActiveTimestamp = System.currentTimeMillis())
        .closeHandler(closeHandler);

    sendAssistantMessage(true);

    log.info("client {} socket connect succeed, sid {} aid {}.", handlerID, storeId, assistantId);
  }

  private void sendAssistantMessage(boolean isLogin) {
    final long now = System.currentTimeMillis();
    vertx.executeBlocking(promise -> {
      final JsonObject header = new JsonObject()
          .put("fn", isLogin ? "LOGIN" : "LOGOUT") // EnumWsFnType.LOGIN LOGOUT
          .put("ts", now);
      final JsonObject message = new JsonObject().put("header", header).put("body", token);
      messageSender.send(connectorId, message.toString());
      promise.complete();
    });
  }

  static final WebSocketFrame PING = WebSocketFrame.pingFrame(Buffer.buffer());

  private long registerHeartbeat(ServerWebSocket websocket) {
    return vertx.setPeriodic(heartbeatStepSeconds.get(), timerId -> {

      websocket.writeFrame(PING)
          .onFailure(throwable -> {
            log.error("{} send ping failed", connectorId, throwable);
            if ("WebSocket is closed".equals(throwable.getMessage())) {
              cancelHeartbeat(timerId);
            }
          });

      if ((System.currentTimeMillis() - lastActiveTimestamp)
          > (heartbeatStepSeconds.get() * heartbeatCheckCount.get())) {
        final String handlerID = websocket.textHandlerID();
        log.warn("{} {} {} 长时间未收到 pong, 关闭连接。", storeId, assistantId, handlerID);
        websocket.close((short) 1000, "长时间未收到 pong, 关闭连接。");
      }
    });
  }

  private void cancelHeartbeat(long timerId) {
    if (vertx.cancelTimer(timerId)) {
      log.info("cancelHeartbeat {} succeed", connectorId);
    }
  }

  private void mdcSetup() {
    MDC.put("sid", storeId);
    MDC.put("aid", assistantId);
    MDC.put("cid", connectorId);
  }

}
